import os
import urllib.request as urllib
from torchvision import datasets


def get_datasets(name: str, root: str, download: bool = True):
    if name == 'cifar_10':
        return get_cifar10(root, download)
    elif name == 'stl_10':
        return get_stl10(root, download)
    else:
        raise Exception("wrong datasets name!")


def get_cifar10(root: str, download: bool):
    DATA_URL = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
    if download:
        download_tar_gz(root, DATA_URL)
    tr_set = datasets.CIFAR10(root, True, download=download)
    vl_set = datasets.CIFAR10(root, False, download=False)
    lb_dict = ["airplane", "automobile", "bird", "cat",
               "deer", "dog", "frog", "horse", "ship", "truck"]
    return tr_set, vl_set, lb_dict


def get_stl10(root: str, download: bool):
    DATA_URL = 'http://ai.stanford.edu/~acoates/stl10/stl10_binary.tar.gz'
    if download:
        download_tar_gz(root, DATA_URL)
    tr_set = datasets.STL10(root, "train", download=True)
    vl_set = datasets.STL10(root, "test")
    lb_dict = ['airplane', 'bird', 'car', 'cat',
               'deer', 'dog', 'horse', 'monkey', 'ship', 'truck']
    return tr_set, vl_set, lb_dict


def download_tar_gz(root: str, data_url):
    dest_directory = root
    filename = data_url.split('/')[-1]
    filepath = os.path.join(dest_directory, filename)
    if not os.path.exists(filepath):
        def _progress(count, block_size, total_size):
            print('\rDownloading %s %.2f%%' % (filename, float(
                count * block_size) / float(total_size) * 100.0), end='\r')
        filepath, _ = urllib.urlretrieve(
            data_url, filepath, reporthook=_progress)
        print('Downloaded', filename)
    else:
        print('Already Downloaded')
